#include "eigen/softmax.h"

namespace  ldl_eigen
{
void SoftMax::forward()
{
    // 计算指数
    Eigen::MatrixXf exp_values = (*m_ptr_input).array().exp();
    
    // 计算每行的和
    Eigen::VectorXf row_sums = exp_values.rowwise().sum();
    
    // 归一化得到softmax概率
    m_output = exp_values.array().colwise() / row_sums.array();
}

void SoftMax::backward() {
    Eigen::VectorXf S = (m_output.array() * (*m_ptr_output_gradient).array()).rowwise().sum();
    m_input_gradient = m_output.array() * ((*m_ptr_output_gradient).array() - S.replicate(1, m_output.cols()).array());
}

}